import torch
import torch.nn as nn
from tqdm import tqdm
import os

from datasets import load_dataset


def evaluate_ppl(model, tokenizer, device, dataset_name='wikitext2') -> float:

    print(f"[INFO] Evaluating PPL on {dataset_name}")
    if dataset_name == 'wikitext2':
        # 下载并加载 WikiText-2 原始版测试集
        data = load_dataset('wikitext', 'wikitext-2-raw-v1', split='test')
    else:
        raise ValueError(f"Unsupported dataset: {dataset_name}")

    # 将整张测试集文本 tokenize，padding=True 和 truncation=True 确保统一长度
    encodings = tokenizer(data['text'], return_tensors='pt', padding=True, truncation=True)
    input_ids = encodings.input_ids.to(device)  # 将 token ids 移动到指定设备（GPU 或 CPU）
    bsz, seqlen = input_ids.shape  # bsz 是批大小，seqlen 是序列长度

    # 定义交叉熵损失函数，忽略标签为 -100 的位置，且将所有 loss 求和
    loss_fct = nn.CrossEntropyLoss(ignore_index=-100, reduction='sum')
    total_loss = 0.0
    total_tokens = 0

    # 在不计算梯度的上下文中进行前向和 loss 计算
    with torch.no_grad():
        # 按照模型能接受的最大序列长度分块处理
        for start in range(0, seqlen - 1, model.config.max_position_embeddings):
            end = min(start + model.config.max_position_embeddings, seqlen)
            batch = input_ids[:, start:end]  # 形状 [bsz, chunk_len]
            labels = batch.clone()
            labels[:, :-1] = -100  # 除最后一个 token 外，其余全部设为 -100（忽略项）

            outputs = model(batch)  # 前向传播
            logits = outputs.logits  # 形状 [bsz, chunk_len, vocab_size]

            # 将 logits 和 labels 都向左/向右平移一位以对齐预测与真实标签
            shift_logits = logits[..., :-1, :].reshape(-1, logits.size(-1))  # [(bsz*(chunk_len-1)), vocab_size]
            shift_labels = labels[..., 1:].reshape(-1)  # [(bsz*(chunk_len-1))]

            # 计算当前分块的交叉熵损失之和
            loss = loss_fct(shift_logits, shift_labels)
            total_loss += loss.item()  # 累加到 total_loss
            # 有效 token 数 = (chunk_len - 1) * bsz
            total_tokens += (end - start - 1) * bsz

    # 计算整体 perplexity = exp(total_loss / total_tokens)
    ppl = torch.exp(torch.tensor(total_loss / total_tokens))
    return ppl.item()/10000  # 返回 Python float


def evaluate_ppl_light(model, tokenizer, device, test_file=None, max_tokens=1000):

    print(f"[INFO] Evaluating PPL on local file: {test_file}")
    if not os.path.isfile(test_file):
        # 如果文件不存在，就报错
        raise FileNotFoundError(f"Test file not found: {test_file}")

    # 逐行读取整个测试文件
    with open(test_file, 'r', encoding='utf-8') as f:
        lines = f.readlines()

    # 将所有行用空格拼成一个长字符串
    full_text = " ".join(lines)
    # tokenize 并返回 PyTorch 张量
    enc = tokenizer(full_text, return_tensors='pt')
    input_ids = enc.input_ids[0][:max_tokens].unsqueeze(0).to(device)

    # 构造标签，将除最后一个位置以外标为 -100，表示忽略
    labels = input_ids.clone()
    labels[:, :-1] = -100

    # 前向传播获取 logits
    with torch.no_grad():
        outputs = model(input_ids)  # [1, T, vocab_size]
        logits = outputs.logits

    # 把 logits 去掉最后一个时间步，并展平为 [ (T-1) , vocab_size ]
    shift_logits = logits[..., :-1, :].reshape(-1, logits.size(-1))
    # 把 labels 右移一位，并展平；只有最后一个位置的标签不为 -100，其余都是 -100
    shift_labels = labels[..., 1:].reshape(-1)

    # 定义交叉熵损失，忽略标签为 -100 的位置，reduction='sum' 表示将所有位置的 loss 累加
    loss_fct = nn.CrossEntropyLoss(ignore_index=-100, reduction='sum')
    loss = loss_fct(shift_logits, shift_labels)  # 计算累加 loss
    # 统计实际参与计算的 token 数（标签不为 -100 的数量）
    total_tokens = (shift_labels != -100).sum()

    # 计算 perplexity = exp(loss / total_tokens)
    ppl = torch.exp(loss / total_tokens)
    return ppl.item()  # 返回 Python float
